Skip to content

dataloader: benchmark + packing-bos-resume strategies - #13

Open
art-test-stack wants to merge 5 commits into
masterfrom
dataloader-benchmark
Open

dataloader: benchmark + packing-bos-resume strategies#13
art-test-stack wants to merge 5 commits into
masterfrom
dataloader-benchmark

Conversation

@art-test-stack

@art-test-stack art-test-stack commented Aug 17, 2026

Copy link
Copy Markdown
Owner

Dataloader packing and benchmark overhaul

Summary

This branch makes token packing explicit, measurable, and resumable. It fixes the default stream loader so tokens and adjacent transitions are preserved across rows and batches, adds a lossless BOS-aligned strategy, retains the legacy destructive best-fit strategy for comparison, and rebuilds the dataloader benchmark around policy-matched correctness and performance measurements.

Packing strategies

Strategy Behavior Trade-off
stream Treats documents as one continuous B*T+1 token stream and carries the final token into the next batch. Default and lossless, but rows do not necessarily start with BOS.
bos_aligned Starts every row with BOS, retains document continuations, and inserts a synthetic BOS when a document spans rows. Lossless for complete batches, with synthetic-BOS overhead.
bos_bestfit_crop Uses the legacy nanochat-style best-fit packer to keep rows full and BOS-aligned. May permanently discard document suffixes and retains approximate resume behavior.

Neither stream nor bos_aligned emits padded or incomplete batches. A finite source stops when it cannot fill the next batch.

BOS alignment is a layout policy, not document isolation: without segment-aware attention masks, tokens can still attend across document boundaries within a row.

Key changes

  • Reworked DistDataLoader to preserve FIFO token order, retain long-document tails, and maintain transitions across row and batch boundaries.
  • Added explicit stream, bos_aligned, and bos_bestfit_crop selection through the public API, configuration, and training CLI.
  • Added exact checkpoint state for stream and bos_aligned, including pending tokens, carry tokens, continuation state, the active strategy, and packing statistics.
  • Corrected shard resume offsets so checkpoints continue from the next unread document instead of replaying the previous chunk.
  • Added PackingStats counters for source-token flow, destructive crops, skipped transitions, synthetic BOS tokens, intentional BOS boundaries, and buffered state.
  • Updated the README with each strategy's guarantees and limitations.

Benchmark changes

The dataloader benchmark now compares implementations only when they implement the same packing policy. It:

  • runs correctness checks before timing;
  • separates flat-stream and destructive BOS-best-fit comparisons;
  • supports on-the-fly and pretokenized inputs;
  • measures loader and device-transfer time without model execution;
  • rotates implementation order across trials and excludes warmup work;
  • reports throughput, latency, source-token utilization, transition coverage, supervision utilization, crop rate, BOS alignment, and buffering;
  • writes reproducible metadata plus JSON, CSV, HTML, and optional plot artifacts.

No benchmark result numbers are committed; results depend on the selected corpus, tokenizer, and device.

Test coverage

New tests cover token and transition preservation, long-document tails, destructive-crop accounting, BOS-aligned packing, deterministic validation, exact checkpoint resume, incomplete-batch rejection, benchmark accounting invariants, and warmup exclusion.

@art-test-stack
art-test-stack requested a balanced review from Copilot August 18, 2026 08:17

This comment was marked as low quality.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated no new comments.

Suppressed comments (3)

scripts/benchmark/dataloaders.py:513

  • read_documents() returns early once limit is reached, skipping the validation that the selected corpus contains text documents. If the first limit rows include non-str values (e.g., pretokenized lists), this will silently pass and fail later in tokenization.
        for row_group in range(parquet.num_row_groups):
            documents.extend(parquet.read_row_group(row_group, columns=[column]).column(0).to_pylist())
            if len(documents) >= limit:
                return documents[:limit]
    if not documents or not all(isinstance(value, str) for value in documents):

src/gpt_lab/data/loader.py:582

  • buffer_size or 1000 treats 0 as falsy and overrides an explicit buffer_size=0. If 0 is invalid, it should be rejected explicitly; otherwise, prefer an is None check so caller intent is preserved.
            tokenizer_batch_size=128,
            device=dist_info["DEVICE"],
            resume_state_dict=resume_state_dict,
            buffer_size=buffer_size or 1000,
            base_path=(Path(datadir) if datadir is not None else DATA_DIR) / name,
            packing_stats=packing_stats,

README.md:205

  • This sentence is misleading: build_dataloader is exposed from the gpt_lab.data module, not “through” the DistDataLoader class. Pointing to the module export helps readers find the API quickly.
The `build_dataloader` function, accessible through [`gpt_lab.data.DistDataLoader`](./src/gpt_lab/data/loader.py), exposes three explicit packing strategies:

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants